{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from typing import Optional, Tuple, Type\n",
    "\n",
    "class PatchEmbed(nn.Module):\n",
    "    \"\"\"\n",
    "    Image to Patch Embedding.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        kernel_size: Tuple[int, int] = (16, 16),\n",
    "        stride: Tuple[int, int] = (16, 16),\n",
    "        padding: Tuple[int, int] = (0, 0),\n",
    "        in_chans: int = 3,\n",
    "        embed_dim: int = 768,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            kernel_size (Tuple): kernel size of the projection layer.\n",
    "            stride (Tuple): stride of the projection layer.\n",
    "            padding (Tuple): padding size of the projection layer.\n",
    "            in_chans (int): Number of input image channels.\n",
    "            embed_dim (int): Patch embedding dimension.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        self.proj = nn.Conv2d(\n",
    "            in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = self.proj(x)\n",
    "        # B C H W -> B H W C\n",
    "        x = x.permute(0, 2, 3, 1)\n",
    "        return x\n",
    "    \n",
    "img = torch.rand((1,3,1024,1024))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64, 64, 768])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# init \n",
    "model = PatchEmbed()\n",
    "model(img).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([31, 96])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 位置编码\n",
    "input_size = [16,16]\n",
    "head_dim = 96\n",
    "\n",
    "torch.zeros(2 * input_size[0] - 1, head_dim).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0, 0, 0],\n",
       "         [1, 1, 1],\n",
       "         [2, 2, 2]]),\n",
       " tensor([[0, 1, 2],\n",
       "         [0, 1, 2],\n",
       "         [0, 1, 2]]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y = torch.meshgrid(torch.arange(3), torch.arange(3))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vlm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
